import numpy as np
import SimpleITK as sitk

__all__ = ['load_mhd_image',
           'normalize_ct_lung_to_unit',
           'get_numpy_from_sitk']

def load_mhd_image(filename):
    """
    simple itk load mhd file
    :param filename:inpute filename
    :return:dict
    """
    with open(filename) as f:
        contents = f.readlines()
        line = [k for k in contents if k.startswith('TransformMatrix')][0]
        transformM = np.array(line.split(' = ')[1].split(' ')).astype('float')
        transformM = np.round(transformM)
        if np.any(transformM != np.array([1, 0, 0, 0, 1, 0, 0, 0, 1])):
            isflip = True
        else:
            isflip = False
    itkimage = sitk.ReadImage(filename)
    numpyOrigin = np.array(list(reversed(itkimage.GetOrigin())))
    numpySpacing = np.array(list(reversed(itkimage.GetSpacing())))
    return {"itkimage": itkimage, "origin": numpyOrigin, "spacing": numpySpacing, "flip": isflip}


def get_numpy_from_sitk(sitkimage:sitk.Image):
    """
    convert sitkimage to numpy array
    :param sitkimage:
    :return:
    """
    return sitk.GetArrayFromImage(sitkimage)

def normalize_ct_to_unit(sitkImage, shift, scale):
    """
    normalize sitkImage
    :param sitkImage: sitkImage
    :return:normalized sitkImage
    """
    sitkImage = sitk.Cast(sitkImage, sitk.sitkFloat64)
    normalized_sitkImage = sitk.ShiftScale(sitkImage, shift=shift, scale=scale)
    img_array = sitk.GetArrayFromImage(normalized_sitkImage)
    img_array[img_array < 0.0] = 0.0
    img_array[img_array > 1.0] = 1.0
    origin = normalized_sitkImage.GetOrigin()
    spacing = normalized_sitkImage.GetSpacing()
    try:
        sitkImage = sitk.GetImageFromArray(img_array)
    except Exception as e:
        print(e)
        return None
    sitkImage.SetOrigin(origin)
    sitkImage.SetSpacing(spacing)

    return sitkImage


# normalize lung from [-1100 200] to [0 1]
# sitk.Normalize() is used
def normalize_ct_lung_to_unit(sitkImage):
    _shift = 1100.0
    _scale = 1.0 / 1300.0
    return normalize_ct_to_unit(sitkImage, _shift, _scale)